import Logger from './logger';
import OcrLlmApi from './ocr-llm-task';
import {
    SDKConfig, TaskStatus, OcrResult, OcrLlmResult
} from './types';

export * from './types';

type TaskStatusRecord = Record<string, typeof TaskStatus[keyof typeof TaskStatus]>;
type TaskResult = OcrResult | OcrLlmResult | null;

export default class OcrLlmSDK {
    public api: OcrLlmApi;

    private _logger: Logger;
    private _taskList: TaskStatusRecord[] = [];
    private _taskResultList: TaskResult[] = [];
    private _intervalId: ReturnType<typeof setInterval> | null = null;
    private _onProcess: ((tasks: TaskStatusRecord[], taskResults: TaskResult[]) => void) | null = null;
    private _pollingInterval = 1000;

    constructor(config: SDKConfig) {
        this.api = new OcrLlmApi(config);
        this._logger = new Logger(config.debug ?? false, config.logLevel ?? "debug")
    }

    private async _createTask(
        createFn: () => Promise<any>
    ): Promise<any> {
        const res = await createFn();
        if (res.code === 200 && this._intervalId) {
            this.appendTask(res.data.task_id);
            this._logger.debug(`Task ${res.data.task_id} appended to list.`)
        }
        this._logger.debug(`Task ${res.data.task_id} appending skipped.`)
        return res;
    }

    createTaskByUrl(imageUrl: string, prompt: string = "") {
        return this._createTask(() => this.api.createTaskByUrl(imageUrl, prompt));
    }

    createTaskByBase64(base64Str: string, prompt: string = "") {
        return this._createTask(() => this.api.createTaskByBase64(base64Str, prompt));
    }

    createTaskByBinary(file: File, prompt: string = "") {
        return this._createTask(() => this.api.createTaskByBinary(file, prompt));
    }

    createByUrl(imageUrl: string, prompt: string = "") {
        return this._createTask(() => this.api.createByUrl(imageUrl, prompt));
    }

    createByBase64(base64Str: string, prompt: string = "") {
        return this._createTask(() => this.api.createByBase64(base64Str, prompt));
    }

    createByBinary(file: File, prompt: string = "") {
        return this._createTask(() => this.api.createByBinary(file, prompt));
    }

    async getTaskResult(taskId: string) {
        const res = await this.api.getTaskResult(taskId);
        if (res.code === 200 && this._intervalId) {
            const index = this._taskList.findIndex(task => Object.keys(task)[0] === taskId);
            if (index !== -1) {
                this._taskList[index][taskId] = res.data.status;
                this._taskResultList[index] = res.data;
            }
        }
        return res;
    }

    startListening(intervalMs = this._pollingInterval) {
        this.stopListening();
        this._pollingInterval = intervalMs;

        this._intervalId = setInterval(async () => {
            const waitingTaskIds = this._taskList
                .filter(task => {
                    const status = task[Object.keys(task)[0]];
                    return status === TaskStatus.PENDING || status === TaskStatus.PROCESSING;
                })
                .map(task => Object.keys(task)[0]);

            if (waitingTaskIds.length === 0) return;

            this._logger.info(`Processing tasks: ${JSON.stringify(waitingTaskIds)}`)
            // 并行获取任务结果
            await Promise.all(waitingTaskIds.map(async taskId => {
                const res = await this.api.getTaskResult(taskId);
                if (res) {
                    const index = this._taskList.findIndex(task => Object.keys(task)[0] === taskId);
                    if (index !== -1) {
                        this._taskList[index][taskId] = res.data.status;
                        this._taskResultList[index] = res.data;
                    }
                }
            }));

            if (this._onProcess) {
                this._onProcess(this._taskList, this._taskResultList);
            }
            this._logger.info(`Refreshing tasks: ${JSON.stringify(this._taskList)}`)
        }, this._pollingInterval);
    }

    stopListening() {
        if (this._intervalId) {
            clearInterval(this._intervalId);
            this._intervalId = null;
        }
        this._taskList = [];
        this._taskResultList = [];
    }

    appendTask(taskId: string) {
        if (!taskId || this._taskList.some(task => Object.keys(task)[0] === taskId)) {
            this._logger.warn(`Task id [${taskId}] is invalid or already in process list.`)
            return;
        }
        this._taskList.push({ [taskId]: TaskStatus.PENDING });
        this._logger.info(`Task [${taskId}] appended to process list.`)
    }

    getTaskList() {
        return [...this._taskList];
    }

    getTaskResultList() {
        return [...this._taskResultList];
    }

    onProcessing(cb: (tasks: TaskStatusRecord[], taskResults: TaskResult[]) => void) {
        this._onProcess = cb;
        this._logger.info(`Function ${cb.name || 'callback'} will receive [tasks] and [taskResults] for params.`)
    }
}
